#!/usr/bin/env python
#-*- coding:utf-8 -*-

from displayer.displayer import PyplotDisplayer
from graph.node_graph import NodeGraph
from graph.pixel_graph import PixelGraph
from planner.astar_planner import AStarPlanner
from planner.prm_planner import PRMPlanner
from planner.rrt_planner import RRTPlanner
from planner.birrt_planner import BiRRTPlanner


def demo_1_astar_on_node_graph(save=False, display=False):
    vertices = {"S":(0.05, 0.5), #0 #含position信息
                "A":(0.3, 0.8), #1
                "B":(0.3, 0.5), #2
                "C":(0.3, 0.2), #3
                "D":(0.6, 0.95), #4
                "E":(0.6, 0.65), #5
                "F":(0.6, 0.4),  #6
                "K":(0.8, 0.2), #7
                "G":(0.99, 0.5),} #8
    arcs = { ('S', 'A'): 9, ('S', 'B'): 6, ('S', 'C'): 8, ('A', 'S'): 9, ('B', 'S'): 6, ('C', 'S'): 8, #弧是有向边，边可以去重
                ('A', 'D'): 7, ('A', 'E'): 9, ('D', 'A'): 7, ('E', 'A'): 9,
                ('B', 'E'): 8, ('B', 'F'): 8, ('E', 'B'): 8, ('F', 'B'): 8,
                ('C', 'K'): 20, ('K', 'C'): 20,
                ('D', 'G'): 16, ('G', 'D'): 16,
                ('E', 'G'): 13, ('G', 'E'): 13, 
                ('F', 'G'): 13, ('F', 'K'): 5, ('G', 'F'): 13, ('K', 'F'): 5, 
                ('K', 'G'): 6, ('G', 'K'): 6 }
    heuristics = {  ("S", "G"): 20, #0
                    ("A", "G"): 15, #1
                    ("B", "G"): 17, #2
                    ("C", "G"): 15, #3
                    ("D", "G"): 11, #4
                    ("E", "G"): 12, #5
                    ("F", "G"): 10, #6
                    ("K", "G"): 5,  #7
                    ("G", "G"): 0, } #8

    node_graph = NodeGraph()
    node_graph.load(vertices, arcs, heuristics)
    displayer = PyplotDisplayer()
    node_graph.draw(displayer, annotate=True, weight=True)
    planner = AStarPlanner()
    planner.plan(node_graph, 'S', 'G')
    short_path = planner.path
    short_path.draw(displayer)
    if save:
        displayer.savefig('./data/atar_search_in_node_graph.png')
    if display:
        displayer.display()
    
def demo_2_astar_on_pixel_graph(save=False, display=False):
    map_file = './data/road_map.png'
    start = (20,15)
    goal = (130, 135)
    pixel_graph = PixelGraph()
    pixel_graph.load(map_file)
    displayer = PyplotDisplayer()
    pixel_graph.draw(displayer)
    planner = AStarPlanner()
    planner.plan(pixel_graph, start, goal)
    short_path = planner.path
    short_path.draw(displayer, trans=True)
    if save:
        displayer.savefig('./data/atar_search_in_road_map.png')
    if display:
        displayer.display()
    
def demo_3_prm_on_pixel_graph(save=False, display=False):
    map_file = './data/road_map.png'
    start = (20,15)
    goal = (130, 135)
    pixel_graph = PixelGraph()
    pixel_graph.load(map_file)
    displayer = PyplotDisplayer()
    pixel_graph.draw(displayer)
    planner = PRMPlanner()
    planner.plan(pixel_graph, start, goal, nodes_num=150, k_neighbours=10)
    road_map = planner.road_map
    road_map.draw(displayer, trans=True, n_size=80)
    if planner.result:
        planner.path.draw(displayer, trans=True)
    if save:
        displayer.savefig('./data/prm_search_in_road_map.png')
    if display:
        displayer.display()

def demo_4_rrt_on_pixel_graph(save=False, display=False):
    map_file = './data/road_map.png'
    start = (20,15)
    goal = (130, 135)
    pixel_graph = PixelGraph()
    pixel_graph.load(map_file)
    displayer = PyplotDisplayer()
    pixel_graph.draw(displayer)
    planner = RRTPlanner()
    planner.plan(pixel_graph, start, goal, random_prob=0.4, step_length=8, max_epoch=10000)
    random_tree = planner.random_tree
    random_tree.draw(displayer, trans=True, n_size=30)
    if planner.result:
        planner.path.draw(displayer, trans=True)
        pass
    else:
        print("can't reach goal")
    if save:
        displayer.savefig('./data/rrt_search_in_road_map.png')
    if display:
        displayer.display()

def demo_5_birrt_on_pixel_graph(save=False, display=False):
    map_file = './data/road_map.png'
    start = (20,15)
    goal = (130, 135)
    pixel_graph = PixelGraph()
    pixel_graph.load(map_file)
    displayer = PyplotDisplayer()
    pixel_graph.draw(displayer)
    planner = BiRRTPlanner()
    planner.plan(pixel_graph, start, goal, step_length=8, max_epoch=5000)
    start_tree, goal_tree = planner.random_trees
    start_tree.draw(displayer, b_color='blue', trans=True, n_size=30)
    goal_tree.draw(displayer, b_color='green', trans=True, n_size=30)
    if planner.result:
        planner.path.draw(displayer, trans=True)
        pass
    else:
        print("can't reach goal")
    if save:
        displayer.savefig('./data/birrt_search_in_road_map.png')
    if display:
        displayer.display()


def main():
    # demo_1_astar_on_node_graph(save=False, display=False)
    # demo_2_astar_on_pixel_graph(save=False, display=True)
    # demo_3_prm_on_pixel_graph(save=False, display=True)
    # demo_4_rrt_on_pixel_graph(save=False, display=True)
    demo_5_birrt_on_pixel_graph(save=True, display=True)
    # demo_6_lattice_on_pixel_graph(save=False, display=False)
    # demo_7_rl_on_pixel_graph(save=False, display=False)


if __name__ == "__main__":
    main()
